import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.metrics import confusion_matrix
import os
from collections import defaultdict
from itertools import chain
from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas


class Plotting:
    def __init__(self, output_dir):
        self.output_dir = output_dir
        self.epoch_losses = []
        self.bleu_scores = []
        self.train_losses = []
        self.val_losses = []
        self.train_accuracies = []
        self.val_accuracies = []
        self.gradient_histograms = defaultdict(list)
        self.train_times = []
        self.learning_rates = []
        self.gradients_dict = {}
        os.makedirs(self.output_dir, exist_ok=True)

    def plot_loss_curve(self):
        """Plot loss curve for training & validation"""
        plt.plot(range(1, len(self.train_losses) + 1), self.train_losses, marker='o', label='Train Loss', color='blue')
        plt.plot(range(1, len(self.val_losses) + 1), self.val_losses, marker='s', label='Validation Loss', color='red')
        plt.title('Training & Validation Loss Curve')
        plt.xlabel('Epoch')
        plt.ylabel('Loss')
        plt.legend()
        plt.grid(True)
        plt.savefig(os.path.join(self.output_dir, 'loss_curve.png'))
        plt.close()

    def plot_bleu_curve(self):
        """Plot BLEU score curve"""
        plt.plot(range(1, len(self.bleu_scores) + 1), self.bleu_scores, marker='o', color='green')
        plt.title('BLEU Score Curve')
        plt.xlabel('Epoch')
        plt.ylabel('BLEU Score')
        plt.grid(True)
        plt.savefig(os.path.join(self.output_dir, 'bleu_curve.png'))
        plt.close()

    def plot_loss_vs_accuracy(self):
        """Plot loss vs accuracy curve"""
        fig, ax1 = plt.subplots()

        ax1.plot(range(1, len(self.train_losses) + 1), self.train_losses, color='b', label='Train Loss')
        ax1.plot(range(1, len(self.val_losses) + 1), self.val_losses, color='orange', label='Validation Loss')
        ax1.set_xlabel('Epoch')
        ax1.set_ylabel('Loss', color='b')
        ax1.tick_params(axis='y', labelcolor='b')

        ax2 = ax1.twinx()
        ax2.plot(range(1, len(self.train_accuracies) + 1), self.train_accuracies, color='r', label='Train Accuracy')
        ax2.plot(range(1, len(self.val_accuracies) + 1), self.val_accuracies, color='g', label='Validation Accuracy')
        ax2.set_ylabel('Accuracy', color='r')
        ax2.tick_params(axis='y', labelcolor='r')

        plt.title('Loss and Accuracy Comparison')
        fig.legend(loc='upper right')
        plt.grid(True)
        plt.savefig(os.path.join(self.output_dir, 'loss_vs_accuracy.png'))
        plt.close()

    def plot_gradient_distributions_combined(self):
        fig, ax = plt.subplots(figsize=(18, 10))  # 更宽一些以容纳长图例

        color_map = plt.get_cmap('tab20')
        max_bins = 0

        for i, (name, grad) in enumerate(self.gradients_dict.items()):
            if grad is not None:
                grad_np = grad.detach().cpu().numpy().flatten()
                if grad_np.size == 0:
                    continue
                color = color_map(i % 20)
                ax.hist(grad_np, bins=100, alpha=0.4, label=name, color=color)
                max_bins = max(max_bins, len(grad_np))

        ax.set_title("Gradient Distributions", fontsize=16)
        ax.set_xlabel("Gradient Value", fontsize=12)
        ax.set_ylabel("Frequency", fontsize=12)
        ax.grid(True)

        # 创建图例并放在图右侧，自动扩展画布大小
        legend = ax.legend(loc='upper left', bbox_to_anchor=(1.01, 1.0), fontsize=8, ncol=1, frameon=False)

        # 用 canvas 计算 legend 大小，并调整图像大小
        fig.canvas = FigureCanvas(fig)
        fig.tight_layout(rect=[0, 0, 0.75, 1])  # 初步压缩主图空间

        # 根据图例大小动态扩展保存区域（防止被裁切）
        fig_width, fig_height = fig.get_size_inches()
        fig.subplots_adjust(right=0.75)
        fig.set_size_inches(fig_width + 4, fig_height)  # 加宽图像来容纳图例
        plt.savefig(os.path.join(self.output_dir, 'gradient_distributions_combined.png'), dpi=300, bbox_extra_artists=(legend,), bbox_inches='tight')
        plt.close(fig)

    def plot_training_time(self):
        """Plot training time per epoch"""
        plt.plot(range(1, len(self.train_times) + 1), self.train_times, marker='o', color='cyan')
        plt.title('Training Time per Epoch')
        plt.xlabel('Epoch')
        plt.ylabel('Time (seconds)')
        plt.grid(True)
        plt.savefig(os.path.join(self.output_dir, 'training_time.png'))
        plt.close()

    def plot_learning_rate(self):
        """Plot learning rate curve"""
        plt.plot(range(1, len(self.learning_rates) + 1), self.learning_rates, marker='o', color='black')
        plt.title('Learning Rate Curve')
        plt.xlabel('Epoch')
        plt.ylabel('Learning Rate')
        plt.grid(True)
        plt.savefig(os.path.join(self.output_dir, 'learning_rate_curve.png'))
        plt.close()

    def plot_confusion_matrix(self, true_labels, predicted_labels, labels):
        """Plot confusion matrix"""
        # 扁平化标签列表
        true_labels_flat = list(chain.from_iterable(true_labels))
        predicted_labels_flat = list(chain.from_iterable(predicted_labels))

        # 确保长度一致
        min_len = min(len(true_labels_flat), len(predicted_labels_flat))
        true_labels_flat = true_labels_flat[:min_len]
        predicted_labels_flat = predicted_labels_flat[:min_len]

        cm = confusion_matrix(true_labels_flat, predicted_labels_flat, labels=list(range(len(labels))))

        plt.figure(figsize=(10, 7))
        sns.heatmap(cm, annot=True, fmt='d', cmap='Blues',
                    xticklabels=labels, yticklabels=labels)
        plt.xlabel('Predicted')
        plt.ylabel('True')
        plt.title('Confusion Matrix')
        plt.savefig(os.path.join(self.output_dir, 'confusion_matrix.png'))
        plt.close()

    def update(self, epoch_losses, bleu_scores, train_losses, val_losses,
               train_accuracies, val_accuracies, train_times, learning_rates):
        """Update metrics after each epoch"""
        self.epoch_losses = epoch_losses
        self.bleu_scores = bleu_scores
        self.train_losses = train_losses
        self.val_losses = val_losses
        self.train_accuracies = train_accuracies
        self.val_accuracies = val_accuracies
        self.train_times = train_times
        self.learning_rates = learning_rates

    def save(self, model, true_labels, predicted_labels, labels):
        """Save all visualizations"""
        self.plot_loss_curve()
        self.plot_bleu_curve()
        self.plot_loss_vs_accuracy()
        # self.plot_gradient_distributions_combined()
        self.plot_training_time()
        self.plot_learning_rate()
        # self.plot_confusion_matrix(true_labels, predicted_labels, labels)
